package main

import (
	"context"
	"log"
	"sync"
	"time"

	"golang.org/x/sync/singleflight"
)

var (
	sg      singleflight.Group        // 全局 Group
	cache   = make(map[string]string) // 缓存
	cacheMu sync.Mutex                //缓存互斥锁
)

func main() {
	testGetDataWithTimeout()
}

func testGetDataWithTimeout() {
	var wg sync.WaitGroup
	keys := []string{"user_123", "user_123", "user_123", "user_123", "user_123", "product_456"} // 包含重复key
	for _, key := range keys {
		wg.Add(1)
		go func(k string) {
			defer wg.Done()
			start := time.Now()
			timeout, cancelFunc := context.WithTimeout(context.Background(), 3*time.Second)
			defer cancelFunc()
			data, err := getDataWithTimeout(timeout, k)
			log.Printf("[%s] Time: %v, Data: %q, Err: %v",
				k, time.Since(start), data, err)
		}(key)
	}

	wg.Wait()
}

func testGetData() {
	var wg sync.WaitGroup
	keys := []string{"user_123", "user_123", "user_123", "user_123", "user_123", "product_456"} // 包含重复key

	for _, key := range keys {
		wg.Add(1)
		go func(k string) {
			defer wg.Done()
			start := time.Now()
			data, err := getData(k)
			log.Printf("[%s] Time: %v, Data: %q, Err: %v",
				k, time.Since(start), data, err)
		}(key)
	}

	wg.Wait()
}

// 带缓存的查询方法
func getData(key string) (string, error) {
	// 先查缓存（使用缓存互斥锁）
	cacheMu.Lock()
	if val, ok := cache[key]; ok {
		cacheMu.Unlock()
		return val, nil
	}
	cacheMu.Unlock()

	// 使用 SingleFlight 防止击穿
	val, err, _ := sg.Do(key, func() (interface{}, error) {
		// 实际查询数据库
		data, err := queryFromDB(key)
		if err != nil {
			return nil, err
		}

		// 更新缓存（使用缓存互斥锁）
		cacheMu.Lock()
		defer cacheMu.Unlock()
		cache[key] = data
		return data, nil
	})

	if err != nil {
		return "", err
	}
	return val.(string), nil
}

func getDataWithTimeout(ctx context.Context, key string) (string, error) {
	// 先查缓存（使用缓存互斥锁）
	cacheMu.Lock()
	if val, ok := cache[key]; ok {
		cacheMu.Unlock()
		return val, nil
	}
	cacheMu.Unlock()
	resultCh := sg.DoChan(key, func() (interface{}, error) {
		// 实际查询数据库
		data, err := queryFromDB(key)
		if err != nil {
			return nil, err
		}

		// 更新缓存（使用缓存互斥锁）
		cacheMu.Lock()
		defer cacheMu.Unlock()
		cache[key] = data
		return data, nil
	})
	select {
	case res := <-resultCh:
		return res.Val.(string), res.Err
	case <-ctx.Done(): // 超时处理
		sg.Forget(key) // 忘记这个key
		return "", ctx.Err()
	}

}

// 模拟数据库查询
func queryFromDB(key string) (string, error) {
	//log.Printf("Querying DB for: %s", key)
	time.Sleep(2 * time.Second) // 模拟耗时操作
	return "data_for_" + key, nil
}
